import gym
from gym.wrappers import RecordVideo


def create_env(mode='train'):
    if mode == 'train':
        env = gym.make('CartPole-v1')
        return env
    else:
        test_env = gym.make('CartPole-v1', render_mode='rgb_array')
        return RecordVideo(test_env, "./dqn_videos", episode_trigger=lambda x: True)
